#!/usr/bin/env python

"""
tf_test_006_linear.py
Linear model test example.
"""

__version__     = "1.0.0"
__author__      = "David Qiu"
__email__       = "david@davidqiu.com"
__website__     = "www.davidqiu.com"
__copyright__   = "Copyright (C) 2017, David Qiu. All rights reserved."


import tensorflow as tf
import numpy as np

import pdb


def generate_data():
  data_x = []
  data_y = []

  expected_W = np.array([[3.0, 4.0],
                         [5.0, 6.0]])
  expected_b = np.array([[0.5],
                         [0.2]])

  sample_size = 200

  for i in range(sample_size):
    sample_x = np.random.sample((2, 1)) * 5.0
    sample_y = expected_W.dot(sample_x) + expected_b
    data_x.append(sample_x)
    data_y.append(sample_y)

  return data_x, data_y


if __name__ == '__main__':
  # data
  data_x, data_y = generate_data()

  # definition: model
  with tf.name_scope('model'):
    x = tf.placeholder(tf.float64, shape=(2, 1), name='x')
    W = tf.Variable(np.zeros((2, 2)), tf.float64, name='W')
    b = tf.Variable(np.zeros((2, 1)), tf.float64, name='b')
    y = tf.add(tf.matmul(W, x), b, name='y')

  # definition: loss
  with tf.name_scope('loss'):
    sample_y = tf.placeholder(tf.float64, shape=(2, 1), name='sample_y')
    loss = tf.reduce_sum(tf.square(y - sample_y), name='loss')
    acc_loss = tf.Variable(0.0, dtype=tf.float64, name='acc_loss')
    reset_acc_loss = tf.assign(acc_loss, 0.0, name='reset_acc_loss')
    update_acc_loss = tf.assign(acc_loss, tf.add(acc_loss, loss), name='update_acc_loss')
    tf.summary.scalar('acc_loss', acc_loss)

  # definition: train
  with tf.name_scope('train'):
    optimizer = tf.train.GradientDescentOptimizer(0.01)
    train = optimizer.minimize(loss)

  # definition: summary
  merge_summaries = tf.summary.merge_all()

  # execution
  with tf.Session() as sess:
    summary_writer = tf.summary.FileWriter('./logs/tf_test_006_linear', sess.graph)

    init = tf.global_variables_initializer()
    sess.run(init)

    epochs = 32
    for ep in range(epochs):
      sess.run(reset_acc_loss)
      for i in range(len(data_x)):
        res = sess.run([loss, train, update_acc_loss], {
          x:        data_x[i],
          sample_y: data_y[i]
        })
        #print('epoch: {}, sample: {}, loss: {}'.format(ep, i, res[0]))
      print('epoch: {}, acc_loss: {}'.format(ep, sess.run(acc_loss)))
      merged_summary = sess.run(merge_summaries)
      summary_writer.add_summary(merged_summary, ep)

    res_W_b = sess.run([W, b])
    print('W:')
    print(res_W_b[0])
    print('b:')
    print(res_W_b[1])


